"""
MySQL数据库查询组件

提供MySQL数据库的查询和操作功能
"""

import asyncio
import json
from typing import Any, Dict, List, Optional
from datetime import datetime, date

import mcp.types as types

from .base import BaseMCPComponent

# MySQL支持（可选）
try:
    import mysql.connector
    from mysql.connector import pooling
    MYSQL_AVAILABLE = True
except ImportError:
    MYSQL_AVAILABLE = False


class MySQLQueryComponent(BaseMCPComponent):
    """MySQL数据库查询组件"""
    
    def __init__(self):
        """初始化MySQL查询组件"""
        self._db_pool = None
        self._db_config: Dict[str, Any] = {}
        super().__init__(name="mysql_query")
    
    def is_available(self) -> bool:
        """检查MySQL驱动是否可用"""
        return MYSQL_AVAILABLE
    
    def init_db_pool(self, host="localhost", port=3306, user="", password="", database="", pool_size=5):
        """初始化数据库连接池"""
        if not MYSQL_AVAILABLE:
            print("❌ MySQL驱动未安装")
            return False
        
        self._db_config = {
            "host": host,
            "port": port,
            "user": user,
            "password": password,
            "database": database
        }
        
        try:
            self._db_pool = pooling.MySQLConnectionPool(
                pool_name="mcp_pool",
                pool_size=pool_size,
                pool_reset_session=True,
                **self._db_config
            )
            print(f"✅ 数据库连接池初始化成功 ({pool_size}个连接)")
            return True
        except Exception as e:
            print(f"❌ 数据库连接池初始化失败: {str(e)}")
            return False
    
    @staticmethod
    def _convert_value(obj: Any) -> Any:
        """转换为可序列化格式"""
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
        elif isinstance(obj, bytes):
            return obj.decode('utf-8', errors='ignore')
        elif isinstance(obj, dict):
            return {k: MySQLQueryComponent._convert_value(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple)):
            return [MySQLQueryComponent._convert_value(item) for item in obj]
        return obj
    
    def _format_rows(self, cursor, rows: List) -> Dict[str, Any]:
        """格式化查询结果"""
        if not cursor.description:
            return {"columns": [], "rows": [], "row_count": 0}
        
        columns = [desc[0] for desc in cursor.description]
        formatted_rows = [
            {col: self._convert_value(val) for col, val in zip(columns, row)}
            for row in rows
        ]
        
        return {
            "columns": columns,
            "rows": formatted_rows,
            "row_count": len(formatted_rows)
        }
    
    def _execute_sql(self, sql: str, params: Optional[List] = None, fetch_all: bool = True) -> Dict[str, Any]:
        """执行SQL语句"""
        if not MYSQL_AVAILABLE:
            return {"success": False, "error": "MySQL驱动未安装"}
        
        if not self._db_pool:
            return {"success": False, "error": "数据库连接池未初始化"}
        
        connection = None
        cursor = None
        
        try:
            connection = self._db_pool.get_connection()
            cursor = connection.cursor()
            
            cursor.execute(sql, params) if params else cursor.execute(sql)
            
            # 判断查询类型
            if sql.strip().upper().startswith(('SELECT', 'SHOW', 'DESCRIBE', 'EXPLAIN')):
                rows = cursor.fetchall() if fetch_all else cursor.fetchmany(100)
                return {
                    "success": True,
                    "data": self._format_rows(cursor, rows)
                }
            else:
                connection.commit()
                return {
                    "success": True,
                    "affected_rows": cursor.rowcount,
                    "message": f"操作成功，影响 {cursor.rowcount} 行"
                }
        except mysql.connector.Error as e:
            if connection:
                connection.rollback()
            return {"success": False, "error": f"数据库错误: {str(e)}"}
        except Exception as e:
            if connection:
                connection.rollback()
            return {"success": False, "error": f"执行错误: {str(e)}"}
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()
    
    def _get_tools_definition(self) -> List[types.Tool]:
        """获取工具定义"""
        if not MYSQL_AVAILABLE:
            return []
        
        return [
            types.Tool(
                name="mysql_query",
                description="执行MySQL查询语句（SELECT、SHOW、DESCRIBE等）",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "sql": {"type": "string", "description": "SQL查询语句"},
                        "params": {"type": "array", "description": "SQL参数（可选）", "items": {"type": "string"}}
                    },
                    "required": ["sql"]
                }
            ),
            types.Tool(
                name="mysql_execute",
                description="执行MySQL数据修改语句（INSERT、UPDATE、DELETE等）",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "sql": {"type": "string", "description": "SQL语句"},
                        "params": {"type": "array", "description": "SQL参数（可选）", "items": {"type": "string"}}
                    },
                    "required": ["sql"]
                }
            ),
            types.Tool(
                name="mysql_test_connection",
                description="测试数据库连接是否正常",
                inputSchema={"type": "object", "properties": {}}
            ),
            types.Tool(
                name="mysql_show_tables",
                description="显示当前数据库的所有表",
                inputSchema={"type": "object", "properties": {}}
            ),
            types.Tool(
                name="mysql_describe_table",
                description="显示指定表的结构",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "table_name": {"type": "string", "description": "表名"}
                    },
                    "required": ["table_name"]
                }
            )
        ]
    
    async def handle_tool_call(self, name: str, arguments: Any) -> List[types.TextContent]:
        """处理工具调用"""
        if not MYSQL_AVAILABLE:
            return [types.TextContent(
                type="text",
                text="✗ MySQL功能不可用，请安装: pip install mysql-connector-python"
            )]
        
        loop = asyncio.get_event_loop()
        
        if name == "mysql_query":
            sql = arguments.get("sql", "")
            if not sql:
                return [types.TextContent(type="text", text="✗ 错误: 未提供SQL语句")]
            
            params = arguments.get("params", None)
            result = await loop.run_in_executor(None, self._execute_sql, sql, params, True)
            
            if result["success"]:
                data = result["data"]
                text = f"✓ 查询成功\n\n列数: {len(data['columns'])}\n行数: {data['row_count']}\n\n"
                if data['row_count'] > 0:
                    text += f"列名: {', '.join(data['columns'])}\n\n结果:\n"
                    text += json.dumps(data['rows'], indent=2, ensure_ascii=False)
                else:
                    text += "查询结果为空"
            else:
                text = f"✗ 查询失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        elif name == "mysql_execute":
            sql = arguments.get("sql", "")
            if not sql:
                return [types.TextContent(type="text", text="✗ 错误: 未提供SQL语句")]
            
            params = arguments.get("params", None)
            result = await loop.run_in_executor(None, self._execute_sql, sql, params, False)
            
            if result["success"]:
                text = f"✓ 执行成功\n\n影响行数: {result['affected_rows']}"
            else:
                text = f"✗ 执行失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        elif name == "mysql_test_connection":
            result = await loop.run_in_executor(None, self._execute_sql, "SELECT 1 as test", None, True)
            
            if result["success"]:
                db_info = f"数据库: {self._db_config.get('database', 'N/A')}\n"
                db_info += f"主机: {self._db_config.get('host', 'N/A')}:{self._db_config.get('port', 'N/A')}"
                text = f"✓ 数据库连接正常\n\n{db_info}"
            else:
                text = f"✗ 数据库连接失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        elif name == "mysql_show_tables":
            result = await loop.run_in_executor(None, self._execute_sql, "SHOW TABLES", None, True)
            
            if result["success"]:
                data = result["data"]
                text = f"✓ 数据库表列表\n\n表数量: {data['row_count']}\n\n"
                if data['row_count'] > 0:
                    tables = [list(row.values())[0] for row in data['rows']]
                    text += "表名:\n" + "\n".join(f"  - {table}" for table in tables)
                else:
                    text += "数据库中没有表"
            else:
                text = f"✗ 查询失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        elif name == "mysql_describe_table":
            table_name = arguments.get("table_name", "")
            if not table_name:
                return [types.TextContent(type="text", text="✗ 错误: 未提供表名")]
            
            result = await loop.run_in_executor(None, self._execute_sql, f"DESCRIBE {table_name}", None, True)
            
            if result["success"]:
                data = result["data"]
                text = f"✓ 表结构: {table_name}\n\n"
                text += json.dumps(data['rows'], indent=2, ensure_ascii=False)
            else:
                text = f"✗ 查询失败\n\n错误: {result['error']}"
            
            return [types.TextContent(type="text", text=text)]
        
        else:
            return [types.TextContent(type="text", text=f"✗ 未知工具: {name}")]
    
    def cleanup(self) -> None:
        """清理资源"""
        super().cleanup()
